import numpy as np
import cv2

from tensorflow.keras.applications import VGG16, VGG19, Xception, ResNet50, \
    InceptionV3, imagenet_utils
from tensorflow.keras.applications.inception_v3 import preprocess_input
from tensorflow.keras.preprocessing.image import img_to_array, load_img

MODELS = {'vgg16': VGG16, 'vgg19': VGG19, 'xception': Xception,
          'inception': InceptionV3, 'resnet': ResNet50}


def error(key: str):
    """
    应用model错误

    :param key: model名称
    :return:
    """
    if not MODELS.get(key, None):
        raise ValueError("没有这个model！")


def main(key: str, pic: str):
    error(key)
    input_shape = (224, 224)
    preprocess = imagenet_utils.preprocess_input
    if key in ("inception", "xception"):
        input_shape = (299, 299)
        preprocess = preprocess_input
    # 加载model
    print(f"[info]:加载{key}模型中...")
    model = MODELS[key](weights="imagenet")

    # 加载和预处理图像
    print("[info]:加载和预处理图像中...")
    image = load_img(pic, target_size=input_shape)
    image = img_to_array(image)
    # 改变形状，将（2， 2， 3）改为（1， 2， 2， 3）以通过网络
    image = np.expand_dims(image, axis=0)
    # 根据已加载的模型（即平均减法、缩放等），使用适当的函数对图像进行预处理
    image = preprocess(image)

    # 开始分类
    print(f"[info]:使用{key}模型进行图像分类中...")
    preds = model.predict(image)
    p = imagenet_utils.decode_predictions(preds)

    # 循环预测结果并向终端显示排名5的预测+概率
    for num, (_, label, prob) in enumerate(p[0], start=1):
        print(f"{num}.{label}--{(prob * 100):.2f}%")

    # 显示图像
    raw = cv2.imread(pic)
    _, label, prob = p[0][0]
    cv2.putText(raw, f"{label=}", (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
    cv2.imshow("Classification", raw)
    cv2.waitKey(0)


if __name__ == '__main__':
    main("vgg16", "D:/bee/dataset/dog_cat/dog.500.jpg")
